package com.example.demo.service.aiService;

import com.example.demo.demo.invoke.DialogueState;
import com.example.demo.enums.ActionType;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

@Service
public class StreamingGenerationService {
    private final ChatClient chatClient;
    private final ObjectMapper objectMapper;
    private static final String GENERATION_PROMPT_TEMPLATE = """
        You are "恋爱大师", a world-class empathetic AI relationship advisor.
        Your designated task for this turn is: **{actionToTake}**.
        
        ### Full Dialogue Context ###
        {currentStateJson}
        
        ### User's Latest Message ###
        "{latestQuery}"
        
        ### Your Task ###
        Based on your designated task and the full context, generate your response to the user now.
        - If your task is `ASK_A_CLARIFYING_QUESTION`, ask one warm, open-ended question.
        - If your task is `PROVIDE_FINAL_ADVICE`, provide the full, three-part structured response: 【心理分析】, 【行动建议】, 【话术参考】.
        """;

    public StreamingGenerationService(ChatClient.Builder chatClientBuilder, ObjectMapper objectMapper) {
        // 初始化基于内存的对话记忆
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(new InMemoryChatMemoryRepository())
                .maxMessages(10)
                .build();
        this.chatClient = chatClientBuilder
                .defaultAdvisors(
                        MessageChatMemoryAdvisor.builder(chatMemory).build()
                )
                .build();
        this.objectMapper = objectMapper;
    }

    public Flux<String> streamReply(DialogueState currentState, ActionType action, String latestQuery) {
        String currentStateJson = convertStateToJson(currentState);
        return this.chatClient.prompt()
                .advisors(advisor -> advisor.param("chat_memory_conversation_id", "678")
                        .param("chat_memory_response_size", 100))
                .user(userSpec -> userSpec
                        .text(GENERATION_PROMPT_TEMPLATE)
                        .param("actionToTake", action.toString())
                        .param("currentStateJson", currentStateJson)
                        .param("latestQuery", latestQuery))
                .stream()
                .content();
    }

    private String convertStateToJson(DialogueState state) {
        try {
            return objectMapper.writeValueAsString(state);
        } catch (JsonProcessingException e) {
            throw new RuntimeException("Error serializing DialogueState to JSON", e);
        }
    }
}